import onnx
import numpy as np

# 加载ONNX模型
onnx_model_path = "/home/lurker/文档/playground/python/pytorch/mnist_model.onnx"
model = onnx.load(onnx_model_path)

# 查看模型的权重（tensor）
# 假设模型的权重在某些特定层中
for tensor in model.graph.initializer:
    print(f"原始张量：{tensor.name}, 数据类型：{tensor.data_type}, 形状：{tensor.dims}")

    # 将每个张量的数值量化为int8
    weight_data = np.frombuffer(tensor.raw_data, dtype=np.float32)
    quantized_data = np.round(weight_data * 127.0).astype(np.int8)  # 量化为int8
    
    # 重新设置量化后的数据
    tensor.raw_data = quantized_data.tobytes()
    tensor.data_type = onnx.TensorProto.INT8  # 设置为int8类型

# 保存量化后的模型
onnx.save(model, "/home/lurker/文档/playground/python/pytorch/quantized_mnist_model.onnx")
print("量化后的模型已保存为 quantized_mnist_model.onnx")
